'''
Created on Feb 19, 2013

@author: nzhao
'''
import numpy as np

from traits.api import Array
from math_qm.QMobj import QMobject
from math_qm.QMoperator import QMoperator

QUBIT_STATES = ['thermal', '+z', '-z', '+x', '-x', '+y', '-y']

class PureState(QMobject):
    state_vect = Array
    
class QMState(QMoperator):
    
    def __init__(self, nbody = 0, dim_list = [], mat = None, entries = [], qubit_state = None):
        if qubit_state in QUBIT_STATES:
            QMobject.__init__(self, 1, [2])
            self.entries = []
            
            if qubit_state is 'thermal':
                self.mat = np.matrix('0.5, 0.0; 0.0, 0.5')
            elif qubit_state is '+x':
                self.mat = np.matrix('0.5, 0.5; 0.5, 0.5')
            elif qubit_state is '-x':
                self.mat = np.matrix('0.5,-0.5;-0.5, 0.5')
            elif qubit_state is '+y':
                self.mat = np.matrix('0.5, 0.5* 1.j; -0.5* 1.j, 0.5')
            elif qubit_state is '-y':
                self.mat = np.matrix('0.5, -0.5* 1.j; 0.5* 1.j, 0.5')
            elif qubit_state is '+z':
                self.mat = np.matrix('1., 0.0; 0.0, 0.0')
            elif qubit_state is '-z':
                self.mat = np.matrix('0.0, 0.0; 0.0, 1.0')
        else:
            QMoperator.__init__(self, nbody, dim_list, mat, entries)
            
if __name__== '__main__':
    state1=QMState(qubit_state='thermal')
    print state1.mat
    print

    state2=QMState(nbody = 1, dim_list=[3], \
                mat = np.matrix('1.0, 0.0, 0.0; \
                                 0.0, 0.0, 0.0; \
                                 0.0, 0.0, 0.0') 
                )
    print state2.mat
    print
    
    state3= QMState(entries=[state1, state2])
    print state3.mat